# drl_policy_driver.py
import torch
import numpy as np
import networkx as nx

import time

import os
import sys

pursuer_dir = os.path.join(os.path.dirname(__file__), "pursuer")
sys.path.append(pursuer_dir)


from pursuer.model import PolicyNet
from pursuer.parameter import INPUT_DIM, EMBEDDING_DIM, K_SIZE, N_ROBOTS

import hashlib  # 新增：用于图哈希缓存


class DRLPolicy:
    def __init__(
        self,
        adj_matrix,
        model_path="./pursuer/model/pursuer_model/policy.pth",
    ):
        self.device = (
            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        )
        self.global_network = PolicyNet(INPUT_DIM, EMBEDDING_DIM).to(self.device)

        # 加载模型权重
        print(f"[{time.strftime('%H:%M:%S')}] Loading model from {model_path}...")
        if self.device == torch.device("cuda"):
            checkpoint = torch.load(f"{model_path}")
        else:
            checkpoint = torch.load(
                f"{model_path}", map_location=torch.device("cpu"), weights_only=True
            )
        print(f"[{time.strftime('%H:%M:%S')}] Model loaded.")

        self.global_network.load_state_dict(checkpoint)
        self.global_network.eval()
        print(f"[{time.strftime('%H:%M:%S')}] Model weights loaded into network.")

        # 存储邻接矩阵
        self.adj_matrix = adj_matrix.copy()
        self.node_num = adj_matrix.shape[0]

        # 新增：图哈希缓存
        self.last_adj_hash = hashlib.md5(adj_matrix.tobytes()).hexdigest()

        # 构建图结构
        self.graph = nx.from_numpy_array(self.adj_matrix)
        print(
            f"[{time.strftime('%H:%M:%S')}] Graph structure created with {self.node_num} nodes."
        )

        # 预计算距离矩阵（优化：使用 all_pairs_shortest_path_length）
        print(f"[{time.strftime('%H:%M:%S')}] Precomputing distance matrix...")
        dist_dict = dict(nx.all_pairs_shortest_path_length(self.graph))
        self.dist_matrix = np.array(
            [
                [
                    dist_dict.get(i, {}).get(j, float("inf"))
                    for j in range(self.node_num)
                ]
                for i in range(self.node_num)
            ]
        )
        # 处理孤立节点（inf → max_dist 作为 fallback）
        if np.any(np.isinf(self.dist_matrix)):
            self.dist_matrix = np.nan_to_num(self.dist_matrix, nan=self.max_dist)
        self.max_dist = np.max(self.dist_matrix) if self.node_num > 0 else 1.0
        print(f"[{time.strftime('%H:%M:%S')}] Distance matrix precomputed.")

        # 预计算邻接列表
        self.edge_inputs = []
        for node in range(self.node_num):
            neighbors = list(self.graph.neighbors(node))
            # 只填充，不截断（假设图的度数 <= K_SIZE）
            while len(neighbors) < K_SIZE:
                neighbors.append(-1)
            self.edge_inputs.append(neighbors)

        # 创建边掩码
        self.edge_mask = (
            torch.from_numpy(self.adj_matrix).float().unsqueeze(0).to(self.device)
        )
        print(f"[{time.strftime('%H:%M:%S')}] Edge mask created.")
        print(f"[{time.strftime('%H:%M:%S')}] DRLPolicy initialization complete.")

    def update_graph(self, adj_matrix):
        """
        更新图结构和相关预计算数据，而不重新加载模型

        Args:
            adj_matrix: 新的邻接矩阵
        """
        print(f"[{time.strftime('%H:%M:%S')}] Updating graph structure...")

        # 新增：哈希检查，跳过无变化更新
        adj_hash = hashlib.md5(adj_matrix.tobytes()).hexdigest()
        if self.last_adj_hash == adj_hash:
            print(f"[{time.strftime('%H:%M:%S')}] Graph unchanged, skipping update.")
            return
        self.last_adj_hash = adj_hash

        self.adj_matrix = adj_matrix.copy()
        self.node_num = adj_matrix.shape[0]

        # 构建新图结构
        self.graph = nx.from_numpy_array(self.adj_matrix)

        print(
            f"[{time.strftime('%H:%M:%S')}] New graph structure created with {self.node_num} nodes."
        )

        # 预计算新的距离矩阵（优化：使用 all_pairs_shortest_path_length）
        print(f"[{time.strftime('%H:%M:%S')}] Updating distance matrix...")
        dist_dict = dict(nx.all_pairs_shortest_path_length(self.graph))
        self.dist_matrix = np.array(
            [
                [
                    dist_dict.get(i, {}).get(j, float("inf"))
                    for j in range(self.node_num)
                ]
                for i in range(self.node_num)
            ]
        )
        # 处理孤立节点（inf → max_dist 作为 fallback）
        if np.any(np.isinf(self.dist_matrix)):
            self.dist_matrix = np.nan_to_num(self.dist_matrix, nan=self.max_dist)
        self.max_dist = np.max(self.dist_matrix) if self.node_num > 0 else 1.0
        print(f"[{time.strftime('%H:%M:%S')}] Distance matrix updated.")

        # 预计算新的邻接列表
        self.edge_inputs = []
        for node in range(self.node_num):
            neighbors = list(self.graph.neighbors(node))
            # 只填充，不截断（假设图的度数 <= K_SIZE）
            while len(neighbors) < K_SIZE:
                neighbors.append(-1)
            self.edge_inputs.append(neighbors)

        print(f"[{time.strftime('%H:%M:%S')}] Edge inputs updated.")

        # 更新边掩码
        self.edge_mask = (
            torch.from_numpy(self.adj_matrix).float().unsqueeze(0).to(self.device)
        )
        print(f"[{time.strftime('%H:%M:%S')}] Edge mask updated.")
        print(f"[{time.strftime('%H:%M:%S')}] DRLPolicy class defined.")

    def calculate_next_move(self, evader_node, pursuer1_node, pursuer2_node):
        """
        基于当前状态计算下一步移动

        Args:
            evader_node: 逃逸者节点索引
            pursuer1_node: 追捕者1节点索引
            pursuer2_node: 追捕者2节点索引

        Returns:
            tuple: (pursuer1_next_node, pursuer2_next_node) 追捕者的下一步节点
        """
        # 构造节点特征
        node_feature = self._construct_node_features(
            evader_node, pursuer1_node, pursuer2_node
        )

        # 为两个追捕者计算下一步动作
        pursuer1_next_node = self._calculate_single_move(node_feature, pursuer1_node, 0)
        pursuer2_next_node = self._calculate_single_move(node_feature, pursuer2_node, 1)

        return pursuer1_next_node, pursuer2_next_node

    def _construct_node_features(self, evader_node, pursuer1_node, pursuer2_node):
        """
        构造节点特征矩阵

        Args:
            evader_node: 逃逸者节点索引
            pursuer1_node: 追捕者1节点索引
            pursuer2_node: 追捕者2节点索引

        Returns:
            torch.Tensor: 节点特征张量 (1, node_num, 3)
        """
        node_features = []

        for i in range(self.node_num):
            # 特征顺序：[到逃逸者的最短路径距离（归一化）, 追捕者1是否存在, 追捕者2是否存在]
            dist_to_evader = self.dist_matrix[i, evader_node] / self.max_dist
            is_p1 = 1.0 if i == pursuer1_node else 0.0
            is_p2 = 1.0 if i == pursuer2_node else 0.0
            feature = [dist_to_evader, is_p1, is_p2]
            node_features.append(feature)

        node_features_tensor = (
            torch.FloatTensor(node_features).unsqueeze(0).to(self.device)
        )
        return node_features_tensor

    def _calculate_single_move(self, node_inputs, current_node_index, robot_index):
        """
        计算单个机器人的下一步移动

        Args:
            node_inputs: 节点特征输入
            current_node_index: 当前节点索引
            robot_index: 机器人索引 (0 或 1)

        Returns:
            int: 下一步节点索引
        """
        # 获取当前节点的邻居
        edge = self.edge_inputs[current_node_index][:]  # 复制列表
        edge_input = torch.tensor(edge).unsqueeze(0).unsqueeze(0).to(self.device)

        # 创建边填充掩码
        edge_padding_mask = torch.zeros((1, 1, K_SIZE), dtype=torch.int64).to(
            self.device
        )
        one = torch.ones_like(edge_padding_mask, dtype=torch.int64).to(self.device)
        edge_padding_mask = torch.where(edge_input == -1, one, edge_padding_mask)
        edge_inputs = torch.where(edge_input == -1, 0, edge_input)

        # 当前节点索引
        current_index = (
            torch.tensor([current_node_index]).unsqueeze(0).unsqueeze(0).to(self.device)
        )

        # 节点填充掩码（无填充）
        node_padding_mask = None

        # 执行推理
        with torch.no_grad():
            logp = self.global_network(
                node_inputs,
                edge_inputs,
                current_index,
                node_padding_mask,
                edge_padding_mask,
                self.edge_mask,
                greedy=True,  # 贪婪策略
            )

        # 选择概率最高的动作
        action_index = torch.argmax(logp, dim=1).long()
        next_node_index = edge_inputs[0, 0, action_index.item()]

        return next_node_index
